import os

def rename_images_and_labels(parent_dir, img_extensions, label_ext):
    """
    对目录下全部数据以及标签重命名，改为{当前目录}_{i}的形式
    """
    # 遍历父目录下的所有子目录
    for subdir, dirs, files in os.walk(parent_dir):
        file_index = 1  # 文件索引，用于命名
        for file in sorted(files):  # 排序文件以保持一致的重命名顺序
            # 检查文件是否是图像文件
            if file.lower().endswith(tuple(img_extensions)):
                # 获取文件的原始扩展名
                original_ext = os.path.splitext(file)[-1]
                # 构造新的文件名（当前目录名_下划线_数字_原扩展名）
                new_name = os.path.basename(subdir) + '_' + str(file_index) + original_ext
                # 构造完整的文件路径
                old_path = os.path.join(subdir, file)
                new_path = os.path.join(subdir, new_name)
                # 重命名图像文件
                os.rename(old_path, new_path)
                print(f"Renamed {old_path} to {new_path}")
                file_index += 1
                
                # 寻找并重命名对应的标签文件
                label_file = file[:-len(original_ext)] + label_ext
                label_path = os.path.join(subdir, label_file)
                if os.path.exists(label_path):
                    new_label_name = os.path.basename(subdir) + '_' + str(file_index - 1) + label_ext
                    new_label_path = os.path.join(subdir, new_label_name)
                    os.rename(label_path, new_label_path)
                    print(f"Renamed {label_path} to {new_label_path}")

if __name__ == "__main__":
    parent_directory = r'/home/hxt/nfs_storage/HITCRT/QuanJu'  # 父目录路径
    image_extensions = ('.jpg', '.png')  # 图像文件的扩展名
    label_extension = '.txt'  # 标签文件的扩展名
    rename_images_and_labels(parent_directory, image_extensions, label_extension)